#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import glob
import time
import multiprocessing
import torch

sys.path.append('.')

"""
    人脸
"""


def show_images(image_dir, file_names, total_image_count, dis_text='', win_name='image', wait=False):
    # h_image_count = 10
    # v_image_count = 6
    # img_width = 128
    # img_height = 128
    # h_image_count = 8
    # v_image_count = 4
    # img_width = 224
    # img_height = 224
    h_image_count = 20
    v_image_count = 10
    img_width = 96
    img_height = 96

    display_image_count = h_image_count * v_image_count
    files = []
    last_idx = len(file_names) - 1
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            file_path = os.path.join(image_dir, file_name.strip())
            img = cv2.imread(file_path)
            img = cv2.resize(img, (img_width, img_height))
            # w, h = img.shape[1], img.shape[0]
            # wscale = img_width / w
            # hscale = img_height / h
            # scale = min(wscale, hscale)
            # w = int(round(w * scale))
            # h = int(round(h * scale))
            # image = cv2.resize(img, (w, h))
            # img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
            # x = (img_width-w)//2
            # y = (img_height-h)//2
            # img[y:y+h, x:x+w, :] = image
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%s %d/%d/%d' % (dis_text, idx + 1, last_idx + 1, total_image_count), (8, 24), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return
        else:
            break


def read_features(file_list):
    faceid_sub_mode_files = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            sub_faceid = lines[1]
            file_name = lines[2]
            mode = int(lines[3])

            sub_mode_files = faceid_sub_mode_files.get(faceid, None)
            if sub_mode_files is None:
                files = [file_name]
                mode_files = {mode: files}
                sub_mode_files = {sub_faceid: mode_files}
                faceid_sub_mode_files[faceid] = sub_mode_files
            else:
                mode_files = sub_mode_files.get(sub_faceid, None)
                if mode_files is None:
                    files = [file_name]
                    mode_files = {mode: files}
                    sub_mode_files[sub_faceid] = mode_files
                else:
                    files = mode_files.get(mode, None)
                    if files is None:
                        files = [file_name]
                        mode_files[mode] = files
                    else:
                        files.append(file_name)

        print('Read Feature Count: {}'.format(file_count))

    return faceid_sub_mode_files


def main():
    parser = argparse.ArgumentParser(description="Display Face ReID Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data/data/FaceImage/Image", help="path to image file list", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/Feature.sel.FR.FR.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--face_id", default="", help="only display this face id", type=str
    )
    parser.add_argument(
        "--sub_face_id", default="", help="only display this sub face id", type=str
    )
    args = parser.parse_args()

    faceid_sub_mode_files = read_features(args.file_list)

    if (args.face_id == '') or (args.sub_face_id == '') or (args.face_id is None) or (args.sub_face_id is None):
        for faceid, sub_mode_files in faceid_sub_mode_files.items():
            for subfaceid, mode_files in sub_mode_files.items():
                for mode, files in mode_files.items():
                    win_name = '{}_{}_{}'.format(faceid, subfaceid, mode)
                    show_images(args.image_dir, files, len(files), win_name=win_name)
                key = cv2.waitKey()
                if key == 27 or key == ord('q') or key == ord('Q'):
                    return
                cv2.destroyAllWindows()
    else:
        faceid, subfaceid = args.face_id, args.sub_face_id
        mode_files = faceid_sub_mode_files[faceid][subfaceid]
        for mode, files in mode_files.items():
            win_name = '{}_{}_{}'.format(faceid, subfaceid, mode)
            show_images(args.image_dir, files, len(files), win_name=win_name, wait=True)

    print('finish!')


if __name__ == '__main__':
    main()
